import numpy as np
import odl
from odl.contrib import torch as odl_torch


def build_gemotry_deeplesion():
    """
    目前锥束加入到网络中都比较难，需要的显存过多。
    因此测试使用deeplesion数据，使用扇形束即可，但是参数和InDuDoNet使用的可能不太一样。
    """
    # reco_space_h = odl.uniform_discr(
    #     min_pt=[-param.param['sx'] / 2.0, -param.param['sy'] / 2.0],
    #     max_pt=[param.param['sx'] / 2.0, param.param['sy'] / 2.0], shape=[param.param['nx_h'], param.param['ny_h']],
    #     dtype='float32')
    # Todo 还没搞懂这个参数是干啥用的，先用官方代码中的示例
    reso = 0.05
    reco_space = odl.uniform_discr(
        min_pt=[-128 * reso, -128 * reso], max_pt=[128 * reso, 128 * reso], shape=[256, 256], dtype='float32')
    # startangle 0 endangle 2Pi nProj 720  # 320来自于ADN中制作deeplesion数据时使用的参数（光源旋转次数）
    angle_partition = odl.uniform_partition(0, 2 * np.pi, 320)
    # su 38.0140605566 nu_h 641  0.171来自于ADN中制作deeplesion数据时matlab函数para2fan源码中gammaRad变量
    detector_partition_h = odl.uniform_partition(-184 * reso, 184 * reso, shape=369)
    # dso 39.69230769230769  dde 39.69230769230769
    geometry_h = odl.tomo.FanBeamGeometry(angle_partition,
                                          detector_partition_h,
                                          # src_to_det_init=(0, 1075),
                                          src_radius=1075 * reso,  # Radius of the source circle
                                          det_radius=0,
                                          )  # Radius of the detector circle
    # print(geometry_h.src_position(0))
    ray_trafo = odl.tomo.RayTransform(reco_space, geometry_h, impl='astra_cuda')
    FBP = odl.tomo.fbp_op(ray_trafo, filter_type='Shepp-Logan', )  # frequency_scaling=0.8)
    return ray_trafo, FBP


def generate_deeplesion_odl_operator():
    """
    实验发现，如果没有滤波的话，重建出来的图像的值域会比原始图的大很多
    """
    project, FBP = build_gemotry_deeplesion()  # fp是一个RayTransform类
    # fp = build_gemotry_own()
    # op_module_project = odl_torch.OperatorModule(project)  # OperatorModule类用来封装
    # op_module_backproj = odl_torch.OperatorModule(project.adjoint)
    return odl_torch.OperatorModule(project), odl_torch.OperatorModule(FBP)
